/*
 * Licensed Materials - Property of tenxcloud.com
 * (C) Copyright 2018 TenxCloud. All Rights Reserved.
 * 2018-06-19  @author lizhen
 */

package utility

import "sync"

func NewPool(size int) *pool {
	return &pool{
		workers:  make(chan func(), size),
		stop:     make(chan interface{}, 1),
		join:     make(chan interface{}, 1),
		stopping: false,
		syncer:   &sync.WaitGroup{},
	}
}

type pool struct {
	workers  chan func()
	stop     chan interface{}
	join     chan interface{}
	stopping bool
	syncer   *sync.WaitGroup
}

func (p *pool) Enqueue(worker func()) {
	if p.stopping {
		panic("stopping no new worker can enqueue")
	}
	p.workers <- worker
}

func (p *pool) run() {
	for {
		select {
		case job := <-p.workers:
			if !p.stopping {
				p.syncer.Add(1)
				go func() {
					defer p.syncer.Done()
					job()
				}()
			}
		case <-p.stop:
			p.stopping = true
			count := len(p.workers)
			if count <= 0 {
				return
			}
			p.syncer.Add(count)
			for i := 0; i < count; i++ {
				j := <-p.workers
				go func() {
					defer p.syncer.Done()
					j()
				}()
			}
			p.syncer.Wait()
			p.join <- struct{}{}
			return
		}
	}
}

func (p *pool) Start() {
	go p.run()
}

func (p *pool) Stop() {
	p.stop <- struct{}{}
}

func (p *pool) Join() {
	<-p.join
}

func (p *pool) Close() {
	close(p.workers)
	close(p.stop)
	close(p.join)
}
